import os
import pathlib
import time
import cv2
import imutils
import shutil
import img2pdf
import glob
import argparse

############# 定义常量

OUTPUT_SLIDES_DIR = f"./output"  # 默认输出文件夹

FRAME_RATE = 3  # 帧率：每秒需要处理的帧数越少，速度越快
WARMUP = FRAME_RATE  # 被跳过的初始帧数
FGBG_HISTORY = FRAME_RATE * 15  # 背景对象中的帧数
VAR_THRESHOLD = 16  # 方差阈值，用于判断当前像素是前景还是背景。一般默认为 16，如果光照变化明显，如阳光下的水面，建议设为 25，值越大灵敏度越低
DETECT_SHADOWS = False  # 是否检测影子，设为 True 为检测，False 为不检测，检测影子会增加程序时间复杂度，一般设置为 False
MIN_PERCENT = 0.1  # 在前景和背景之间的最小差值百分比，以检测运动是否已经停止
MAX_PERCENT = 3  # 在前景和背景之间的最大百分比的差异，以检测帧是否仍在运动


def get_frames(video_path):
    '''从位于 video_path 的视频返回帧的函数
    此函数跳过 FRAME_RATE 中定义的帧'''

    # 打开指向视频文件的指针初始化帧的宽度和高度
    vs = cv2.VideoCapture(video_path)
    if not vs.isOpened():
        raise Exception(f'unable to open file {video_path}')

    # 获取视频总帧数
    total_frames = vs.get(cv2.CAP_PROP_FRAME_COUNT)
    # 获取视频帧频率
    fps = vs.get(cv2.CAP_PROP_FPS)
    frame_time = 0
    frame_count = 0
    total_time = 0
    total_time = total_frames / fps
    print("total_frames: ", total_frames)
    print("FRAME_RATE", FRAME_RATE)
    print("Total_time:", total_time)

    # 循环播放视频的帧
    while True:
        # 从视频中抓取一帧

        # 视频帧超过视频总长，则结束
        if(frame_time > total_time):
            break
        vs.set(cv2.CAP_PROP_POS_MSEC, frame_time * 1000)  # 将帧移动到时间戳
        frame_time += 1 / FRAME_RATE

        (_, frame) = vs.read()
        # 如果帧为None，那么我们已经到了视频文件的末尾
        if frame is None:
            break

        frame_count += 1
        yield frame_count, frame_time, frame

    vs.release()


def detect_unique_screenshots(video_path, output_folder_screenshot_path):
    ''''''
    # 用参数初始化 fgbg 一个背景对象
    # history = 影响背景减法器的帧历史数
    # varThreshold = 像素与模型之间的平方马氏距离的阈值，以决定像素是否被背景模型很好地描述。该参数不影响后台更新。
    # detectShadows = 如果为真，算法将检测阴影并标记它们。它会稍微降低速度，因此如果您不需要此功能，请将参数设置为 false。

    fgbg = cv2.createBackgroundSubtractorMOG2(history=FGBG_HISTORY, varThreshold=VAR_THRESHOLD,
                                              detectShadows=DETECT_SHADOWS)

    captured = False
    start_time = time.time()
    (W, H) = (None, None)

    screenshoots_count = 0
    for frame_count, frame_time, frame in get_frames(video_path):
        orig = frame.copy()  # clone the original frame (so we can save it later),
        frame = imutils.resize(frame, width=600)  # resize the frame
        mask = fgbg.apply(frame)  # apply the background subtractor

        # apply a series of erosions and dilations to eliminate noise
        #            eroded_mask = cv2.erode(mask, None, iterations=2)
        #            mask = cv2.dilate(mask, None, iterations=2)

        # if the width and height are empty, grab the spatial dimensions
        if W is None or H is None:
            (H, W) = mask.shape[:2]

        # compute the percentage of the mask that is "foreground"
        p_diff = (cv2.countNonZero(mask) / float(W * H)) * 100

        # if p_diff less than N% then motion has stopped, thus capture the frame

        if p_diff < MIN_PERCENT and not captured and frame_count > WARMUP:
            captured = True
            filename = f"{screenshoots_count:03}_{round(frame_time / 60, 2)}.png"

            path = str(pathlib.Path(output_folder_screenshot_path, filename))

            print("saving {}".format(path))
            cv2.imencode('.png', orig)[1].tofile(path)  # 防止imwrite中文乱码
            screenshoots_count += 1

        # otherwise, either the scene is changing or we're still in warmup
        # mode so let's wait until the scene has settled or we're finished
        # building the background model
        elif captured and p_diff >= MAX_PERCENT:
            captured = False
    print(f'{screenshoots_count} screenshots Captured!')
    print(f'Time taken {time.time() - start_time}s')
    return


def initialize_output_folder(video_path):
    '''Clean the output folder if already exists'''
    (filesname, extension) = os.path.splitext(video_path)
    output_folder_screenshot_path = f"{OUTPUT_SLIDES_DIR}/{video_path.rsplit('/')[-1].replace(extension, '')}"

    if os.path.exists(output_folder_screenshot_path):
        shutil.rmtree(output_folder_screenshot_path)

    os.makedirs(output_folder_screenshot_path, exist_ok=True)
    print('initialized output folder', output_folder_screenshot_path)
    return output_folder_screenshot_path


def convert_screenshots_to_pdf(output_folder_screenshot_path):
    (filesname, extension) = os.path.splitext(video_path)
    output_pdf_path = f"{OUTPUT_SLIDES_DIR}/{video_path.rsplit('/')[-1].replace(extension, '')}" + '.pdf'
    print('output_folder_screenshot_path', output_folder_screenshot_path)
    print('output_pdf_path', output_pdf_path)
    print('converting images to pdf..')
    with open(output_pdf_path, "wb") as f:
        f.write(img2pdf.convert(sorted(glob.glob(f"{output_folder_screenshot_path}/*.png"))))
    print('Pdf Created!')
    print('pdf saved at', output_pdf_path)


if __name__ == "__main__":

    input_dir_name = '/Users/litao/百度网盘/视频'#要转换的视频所在的文件夹

    for file_name in os.listdir(input_dir_name):
        if file_name.endswith(('.mp4', '.avi', '.mov', '.mkv')):
            print(f"正在转换：{file_name}")
            video_path = str(pathlib.Path(input_dir_name, file_name))
            output_folder_screenshot_path = initialize_output_folder(video_path)
            detect_unique_screenshots(video_path, output_folder_screenshot_path)
    
            # 提取的图片转换为pdf
            convert_screenshots_to_pdf(output_folder_screenshot_path)